package com.wcs.exam.common.base;

import cn.hutool.core.util.StrUtil;
import cn.hutool.http.ContentType;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.wcs.exam.common.llm.LlmBaseResp;
import com.wcs.exam.common.llm.LlmFace;
import com.wcs.exam.common.llm.impl.bean.Message;
import lombok.NonNull;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.IOException;
import java.util.List;


/**
 * sse请求工具类
 */
@Slf4j
public class BaseSse {

    @SneakyThrows
    public static SseEmitter stream(@NonNull LlmFace llmFace, @NonNull String requestUrl, @NonNull Headers headers, @NonNull String requestBody) {
        SseEmitter sseEmitter = new BaseSseEmitter();
        Request request = new Request.Builder().url(requestUrl).headers(headers).post(RequestBody.create(requestBody, MediaType.get(ContentType.JSON.getValue()))).build();
        new OkHttpClient().newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(@NonNull Call call, @NonNull IOException e) {
                log.error("请求失败：{}", requestUrl, e);
                sseEmitter.completeWithError(e);
            }

            @Override
            public void onResponse(@NonNull Call call, @NonNull Response response) throws IOException {
                if (response.isSuccessful() && response.body() != null) {
                    BufferedReader reader = new BufferedReader(response.body().charStream());
                    String line;
                    while ((line = reader.readLine()) != null) {
                        if (line.startsWith("data:")) {
                            // 去掉 "data:" 并去除前后的空格
                            line = llmFace.chatHandle(line.replace("data:", "").trim());
                            sseEmitter.send(SseEmitter.event().data(line));
                        }
                    }
                } else {
                    log.error("{}", response);
                }
                response.close();
                sseEmitter.complete();
            }
        });
        return sseEmitter;
    }


    /**
     * 同步请求
     *
     * @param requestUrl  请求地址
     * @param authKey     授权Key
     * @param model       模型
     * @param messageList 消息列表
     * @return BaseResp
     */
    public static LlmBaseResp http(String requestUrl, String authKey, String model, List<Message> messageList) {
        Headers headers = new Headers.Builder().add("Authorization", "Bearer " + authKey).build();
        String requestBody = new JSONObject().putOpt("model", model).putOpt("messages", messageList).putOpt("stream", false).toString();
        Request request = new Request.Builder().url(requestUrl).headers(headers).post(RequestBody.create(requestBody, MediaType.get(ContentType.JSON.getValue()))).build();
        try (Response response = new OkHttpClient().newCall(request).execute()) {
            if (response.isSuccessful() && response.body() != null) {
                String responseBody = response.body().string();
                log.info("llm http请求返回参数：{}", responseBody);

                // 解析响应
                LlmBaseResp baseResp = JSONUtil.toBean(responseBody, LlmBaseResp.class);
                if (StrUtil.isNotBlank(responseBody)) {
                    JSONObject respJson = JSONUtil.parseObj(responseBody);
                    JSONObject resultJson = respJson.getJSONArray("choices").getJSONObject(0);
                    baseResp.setOutput(resultJson.getJSONObject("message").getStr("content"));
                }

                return baseResp;
            } else {
                log.error("请求失败：{}", response);
                throw new IOException("请求失败: " + response);
            }
        } catch (Exception e) {
            log.error("请求异常：{}", e.getMessage(), e);
            return null;
        }
    }
}
